# 这是一个开头
# 人员：Mr Su
# 开发时间：16/2/2021下午12:53
# 文件名：functions.py
# 开发工具：PyCharm
import numpy as np
import matplotlib.pyplot as plt

def MaxMinNormalization(data):
    Max = np.max(data)
    Min = np.min(data)
    data = (data - Min) / (Max - Min)
    return data

def split_train_test(data, test_ratio):
    np.random.seed(123)
    shuffled_indices = np.random.permutation(len(data))
    test_set_size = int(len(data) * test_ratio)
    test_indices = shuffled_indices[:test_set_size]
    train_indices = shuffled_indices[test_set_size:]
    return data[train_indices], data[test_indices]

def data_loader(path):

    datasets = np.load(path)
    # 归一化操作
    datasets = MaxMinNormalization(datasets)

    train_datasets, test_datasets = split_train_test(datasets, test_ratio=0.2)

    return train_datasets,test_datasets

def plot_vae_training_plot(train_losses, test_losses, title, fname):

    # plot_vae_training_plot(train_losses, test_losses, f'Q1({part}) Dataset {dset_id} Train Plot',
    #                        f'results/q1_{part}_dset{dset_id}_train_plot.png')
    # 改动前
    elbo_train, recon_train, kl_train = train_losses[:, 0], train_losses[:, 1], train_losses[:, 2]
    elbo_test, recon_test, kl_test = test_losses[:, 0], test_losses[:, 1], test_losses[:, 2]
    plt.figure()
    #9=n_epochs =9
    n_epochs = len(test_losses) - 1
    #等分区间，分为len(train_losses)份
    x_train = np.linspace(0, n_epochs, len(train_losses))
    # x_test = 1：10
    x_test = np.arange(n_epochs + 1)
    #plot(横坐标，纵坐标，lable)
    plt.plot(x_train, elbo_train, label='-elbo_train')
    plt.plot(x_train, recon_train, label='recon_loss_train')
    plt.plot(x_train, kl_train, label='kl_loss_train')
    plt.plot(x_test, elbo_test, label='-elbo_test')
    plt.plot(x_test, recon_test, label='recon_loss_test')
    plt.plot(x_test, kl_test, label='kl_loss_test')

    plt.legend()
    plt.title(title)
    plt.xlabel('Epoch')
    plt.ylabel('Loss')
    plt.savefig(fname)
    plt.show()